import argparse
import pandas as pd
import numpy as np

from sklearn import metrics, preprocessing
from sklearn import linear_model, model_selection, ensemble

from config import KFOLD_N
import model_dispatcher

import warnings
warnings.filterwarnings('ignore')


def run(fold, traning_model):
    df = pd.read_csv("../input/train_folds.csv")
    test = pd.read_csv("../input/test_clear.csv")
    features = [f for f in df.columns 
                if f not in ('PassengerId', 'Transported', 'kfold', 'Name')]
    one_hot = preprocessing.OneHotEncoder()

    # 处理
    df_train = df[df.kfold != fold].reset_index(drop=True)
    df_valid = df[df.kfold == fold].reset_index(drop=True)
    test = test.reset_index(drop=True)

    full_data = pd.concat(
        [df_train[features], df_valid[features], test[features]],
        axis=0
    )

    one_hot.fit(full_data[features])

    x_train = one_hot.transform(df_train[features])
    x_valid = one_hot.transform(df_valid[features])
    x_test = one_hot.transform(test[features])
    y_train = df_train['Transported']
    y_valid = df_valid['Transported']


    model = model_dispatcher.CLFS[traning_model]
    model.fit(x_train, y_train)
    y_valid_pred = model.predict(x_valid)
    acc = metrics.accuracy_score(y_valid, y_valid_pred)
    acc = round(acc, 6)
    print(f"Fold = {fold}, accuracy = {acc}")

    y_preds = model.predict(x_test)
    result = pd.DataFrame({
        'PassengerId': test.PassengerId, 
        'Transported': y_preds 
    })
    result.to_csv('../output/submission.csv', index=False)
    print("result is in output")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str)
    args = parser.parse_args()
    # for i in range(KFOLD_N): run(i, args.model)
    run(7, args.model)


## 找到一个accuracy 最高的模型
## 输出结果文件，上传到kaggle
## 优化模型